import torch
from torch import nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class Myloss(nn.Module):
    # nn.CrossEntropyLoss

    def __init__(self):
        super(Myloss, self).__init__()

    # pred 为传入的预测值
    # eps 过大将会导致loss出现负值
    def forward(self, pred, label, eps=1e-5):
        # label_shape: tenso([pic1_class,pic2_class])
        batch_size = pred.shape[0]  # 预测值的批次
        # 将输出预测张量转换为[batch, class_num]
        # [2, 3, 1, 1]->[2,3]
        new_pred = pred.reshape(batch_size, -1)
        # 建立一个新张量
        expand_target = torch.zeros(new_pred.shape, device=device)
        for i in range(batch_size):
            # 给图片类别对应位置赋值1
            # one-hot encoding
            expand_target[i, int(label[i])] = 1
        # 将输出预测进行softmax: 按照dim加和求均值后e^x_i/(e^x_1+e^x_2+...)
        # softmax后 0<y^<1, Nan问题暂时不考虑
        softmax_pred = torch.softmax(new_pred, dim=1)
        # 计算总损失 -log(y^{hat}), eps防止梯度爆炸 概率为零时log后负无穷
        return torch.sum(-torch.log(softmax_pred + eps) * expand_target) / batch_size


if __name__ == '__main__':
    import numpy as np

    a = torch.tensor(np.random.randn(2, 5), device=device)
    b = torch.tensor([[0], [1]], device=device)
    print(a.shape)
    loss = Myloss()
    loss1 = torch.nn.CrossEntropyLoss()
    my_loss = loss(a, b)

    print(my_loss)#, my_loss1)
"""


For nn.CrossEntropyLoss the target has to be a single number 
from the interval [0, #classes] instead of a one-hot encoded target vector.
Your target is [1, 0], 
thus PyTorch thinks you want to have multiple labels per input 
which is not supported.

Replace your one-hot-encoded targets:

[1, 0] --> 0

[0, 1] --> 1

"""